import json
import time

from lagent.agents.aggregator import InternLMToolAggregator
from lagent.agents.stream import PLUGIN_CN, AgentForInternLM, MathCoder, get_plugin_prompt
from lagent.llms import INTERNLM2_META, LMDeployPipeline
from lagent.prompts.parsers import PluginParser

# 初始化 LMDeploy 模型
model = LMDeployPipeline(
    path='/root/autodl-tmp/models/internlm2_5-1_8b-chat',  # 模型路径
    meta_template=INTERNLM2_META,  # 使用 InternLM2 的元模板
    tp=1,  # 张量并行度
    top_k=1,  # 只保留概率最高的候选词
    temperature=1.0,  # 采样温度
    stop_words=['<|im_end|>', '<|action_end|>'],  # 停止词
    max_new_tokens=1024,  # 最大生成长度
)

# ----------------------- 测试数学解题能力 -----------------------
print('-' * 80, 'interpreter', '-' * 80)
# 初始化数学解题代理
coder = MathCoder(
    llm=model,
    aggregator=InternLMToolAggregator(few_shot=[[
        # 示例问题：计算邮票价值
        {
            'role': 'user',
            'content': "Mike paid $\\$1.25$ for a stamp three years ago..."  # 问题描述
        },
        {
            'role': 'language',
            'content': "Let's break down the problem..."  # 问题分析
        },
        {
            'role': 'tool',
            'content': '```python\nfrom sympy import symbols...',  # Python代码
            'name': 'interpreter',
        },
        {
            'role': 'environment',
            'content': '20.0',  # 代码执行结果
            'name': 'interpreter'
        },
        {
            'role': 'language',
            'content': 'After calculating...'  # 最终答案
        },
    ]]))

# 测试问题：求解数字问题
query = (
    'Marie is thinking of a multiple of 63, while Jay is thinking of a factor '
    'of 63. They happen to be thinking of the same number. There are two '
    'possibilities for the number that each of them is thinking of, one '
    'positive and one negative. Find the product of these two numbers.')

# 执行求解
res = coder(query, session_id=0)

# 打印解题步骤
print(coder.get_steps(0))

# 打印聚合后的结果
print(
    json.dumps(
        coder.agent.aggregator.aggregate(coder.agent.memory, coder.agent.name,
                                         coder.agent.output_format),
        ensure_ascii=False,
        indent=4))

# 测试状态保存和加载
state_dict = {
    'memory': [
        # 用户输入的问题
        {
            'sender': 'user',
            'content': 'Marie is thinking of a multiple of 63...',
            'formatted': None,
            'type': None,
            'receiver': None
        },
        # 代理的最终答案
        {
            'sender': 'Agent',
            'content': 'The product of the two numbers...',
            'formatted': {
                'tool_type': None,
                'thought': 'The product of the two numbers...',
                'action': None,
                'status': 0
            },
            'type': None,
            'receiver': None
        }
    ],
    'agent.memory': [
        # 详细的解题过程记录
        # ...（包含多个步骤的详细信息）
    ]
}

# 加载状态并验证
coder.load_state_dict(state_dict)
print(coder.state_dict() == state_dict)
print(coder.state_dict())

# ----------------------- 测试插件功能 -----------------------
print('-' * 80, 'plugin', '-' * 80)
# 配置 Arxiv 搜索插件
plugins = [dict(type='lagent.actions.ArxivSearch')]
# 初始化带插件的代理
agent = AgentForInternLM(
    llm=model,
    plugins=plugins,
    output_format=dict(
        type=PluginParser,
        template=PLUGIN_CN,
        prompt=get_plugin_prompt(plugins)))  # 生成插件提示信息

# 测试论文搜索功能
tic = time.time()  # 记录开始时间
query = 'LLM智能体方向的最新论文有哪些？'
res = agent(query, session_id=0)

# 打印结果和执行步骤
print(res.model_dump_json())
print('-' * 120)
print(agent.get_steps(0))
print(f'time elapsed: {time.time() - tic}')  # 打印执行时间
